import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import ssl

# 处理SSL证书问题
ssl._create_default_https_context = ssl._create_unverified_context

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 超参数（需要与训练时保持一致）
input_size = 784
hidden_size = 128
num_classes = 10

# 定义与训练时相同的神经网络模型
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc3(out)
        return out

def load_model(model_path='mnist_model.pth'):
    """加载训练好的模型"""
    model = NeuralNet(input_size, hidden_size, num_classes).to(device)
    try:
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        print(f"成功加载模型: {model_path}")
        return model
    except FileNotFoundError:
        print(f"未找到模型文件 '{model_path}'，请先运行训练脚本")
        return None

def predict(model, image_tensor):
    """对单个图像进行预测"""
    with torch.no_grad():
        image_tensor = image_tensor.reshape(-1, 28*28).to(device)
        output = model(image_tensor)
        predicted = torch.argmax(output).item()
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        confidence = probabilities[predicted].item()
    return predicted, confidence

def test_with_mnist_samples():
    """使用MNIST测试集样本进行推理演示"""
    # 加载模型
    model = load_model()
    if model is None:
        return
    
    # 数据预处理
    transform = transforms.Compose([transforms.ToTensor(), 
                                    transforms.Normalize((0.1307,), (0.3081,))])
    
    # 加载测试数据
    test_dataset = torchvision.datasets.MNIST(root='./data',
                                              train=False,
                                              transform=transform)
    
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=1,
                                              shuffle=True)
    
    # 获取几个测试样本进行预测
    dataiter = iter(test_loader)
    correct = 0
    total = 5  # 测试5个样本
    
    print("=== 推理结果 ===")
    for i in range(total):
        image, true_label = next(dataiter)
        predicted, confidence = predict(model, image)
        
        is_correct = "✓" if predicted == true_label.item() else "✗"
        print(f"样本 {i+1}: 真实={true_label.item()}, 预测={predicted}, 置信度={confidence:.4f} {is_correct}")
        
        if predicted == true_label.item():
            correct += 1
    
    print(f"\n准确率: {correct}/{total} ({100 * correct / total:.1f}%)")

if __name__ == "__main__":
    test_with_mnist_samples()